package com.nlp.mallet;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.io.Reader;
import java.util.ArrayList;
import java.util.Formatter;
import java.util.Iterator;
import java.util.Locale;
import java.util.TreeSet;
import java.util.regex.Pattern;

import cc.mallet.pipe.CharSequence2TokenSequence;
import cc.mallet.pipe.CharSequenceLowercase;
import cc.mallet.pipe.Input2CharSequence;
import cc.mallet.pipe.Pipe;
import cc.mallet.pipe.SerialPipes;
import cc.mallet.pipe.TokenSequence2FeatureSequence;
import cc.mallet.pipe.TokenSequenceRemoveStopwords;
import cc.mallet.pipe.iterator.CsvIterator;
import cc.mallet.topics.ParallelTopicModel;
import cc.mallet.topics.TopicInferencer;
import cc.mallet.types.Alphabet;
import cc.mallet.types.FeatureSequence;
import cc.mallet.types.IDSorter;
import cc.mallet.types.Instance;
import cc.mallet.types.InstanceList;
import cc.mallet.types.LabelSequence;

/**
 * LDA主题模型
 * 在此示例中，我从文件导入数据，训练主题模型，并分析第一个实例的主题分配。
 * 然后，我创建一个新实例，该实例由主题0中的单词组成，并推断该实例的主题分布
 * @author ygsong.abcft
 *
 */
public class TopicModel {
	//static final String DATA_PATH = "E:\\search\\nlp\\machine-learn\\src\\main\\resources\\ap.txt";
	static final String DATA_PATH = "D:\\Mallet\\test\\b.txt";

	public static void main(String[] args) throws IOException {
		ArrayList<Pipe> pipeList = new ArrayList<Pipe>();
		//转小写、过滤标记、去除停用词、映射到特征
		pipeList.add(new Input2CharSequence("UTF-8"));

		pipeList.add(new CharSequenceLowercase());
		pipeList.add( new CharSequence2TokenSequence(Pattern.compile("\\p{L}[\\p{L}\\p{P}]+\\p{L}")) );
        pipeList.add( new TokenSequenceRemoveStopwords(new File("D:\\crf\\data\\dictionary\\stopwords.txt"), "UTF-8", false, false, false) );
        pipeList.add( new TokenSequence2FeatureSequence() );
        SerialPipes serialPipes = new SerialPipes(pipeList);
        
        InstanceList instances = new InstanceList(serialPipes);
        
        Reader fileReader = new InputStreamReader(new FileInputStream(new File(DATA_PATH)),"UTF-8");
       //data, label, name fields
        instances.addThruPipe(new CsvIterator (fileReader, Pattern.compile("^(\\S*)[\\s,]*(\\S*)[\\s,]*(.*)$"),
                3, 2, 1)); // 
        int numTopices = 100;
        //创建一个含有100的topic的model，alpha_t = 0.01, beta_w = 0.01
        //第一个参数作为主题的总和传递
        ParallelTopicModel model = new ParallelTopicModel(numTopices, 1.0, 0.01);
        model.addInstances(instances);
        //使用两个并行采样器，每个查看一半语料库，并在每次迭代后合并统计
        model.setNumThreads(2);
        //运行模型50次迭代并停止（这是为了测试，对于实际应用程序，使用1000到2000次迭代）
        model.setNumIterations(50);
        model.estimate();
        
       
       // 显示第一个实例中的 单词和主题
        
        //数据字母表将单词ID映射到字符串
        Alphabet dataAlpabet = instances.getDataAlphabet();
        
        FeatureSequence tokens = (FeatureSequence) model.getData().get(0).instance.getData();
        LabelSequence topics = model.getData().get(0).topicSequence;
        
        Formatter out = null;
        
        for (int position = 0; position < tokens.getLength(); position++) {
        	out = new Formatter(new StringBuilder(), Locale.CHINA);
			out.format("%s-%d", dataAlpabet.lookupObject(tokens.getIndexAtPosition(position)), topics.getIndexAtPosition(position));
			//System.out.println(out);
        }
        
        System.out.println("******************************************************");
        
        //给定当前状态，预测第一实例的主题分布。
        double[] topicDistribution = model.getTopicProbabilities(0);
        
        //获取一组排序的单词ID
        ArrayList<TreeSet<IDSorter>> topicStoredWords = model.getSortedWords();
        //在第一个文档中显示前5个单词
        for (int topic = 0; topic < numTopices; topic++) {
			Iterator<IDSorter> iterator = topicStoredWords.get(topic).iterator();
			out = new Formatter(new StringBuilder(), Locale.CHINA);
			out.format("%d\t%.3f\t", topic, topicDistribution[topic]);
			int rank = 0;
			while(iterator.hasNext() && rank < 5) {
				IDSorter idCountPair = iterator.next();
				out.format("%s (%.0f) ", dataAlpabet.lookupObject(idCountPair.getID()), idCountPair.getWeight());
				rank++;
			}
			System.out.println(out);
		}
        //从topic_0中创建一个主题概率高的实例
        StringBuilder topicZeroText = new StringBuilder();
        Iterator<IDSorter> iterator = topicStoredWords.get(0).iterator();
        
        int rank = 0;
        while(iterator.hasNext() && rank < 5) {
        	IDSorter idCountPair = iterator.next();
        	topicZeroText.append(dataAlpabet.lookupObject(idCountPair.getID())+" ");
        	rank++;
        }
        //创建一个测试实例
        InstanceList testing = new InstanceList(instances.getPipe());
        testing.addThruPipe(new Instance(topicZeroText.toString(), null, "test instance", null));
        
        TopicInferencer inferencer = model.getInferencer();
        double[] testProbabilities = inferencer.getSampledDistribution(testing.get(0), 10, 1, 5);
        System.out.println("0 \t"+testProbabilities[0]);
	}

}
